package com.jungle.base.dao.impl;

import com.jungle.base.dao.ExtendedJpaRepository;
import jakarta.persistence.EntityManager;
import jakarta.persistence.TypedQuery;
import jakarta.persistence.criteria.CriteriaBuilder;
import jakarta.persistence.criteria.CriteriaDelete;
import jakarta.persistence.criteria.CriteriaQuery;
import jakarta.persistence.criteria.CriteriaUpdate;
import org.springframework.data.domain.Limit;
import org.springframework.data.domain.Sort;
import org.springframework.data.jpa.domain.Specification;
import org.springframework.data.jpa.repository.support.JpaEntityInformation;
import org.springframework.data.jpa.repository.support.SimpleJpaRepository;
import org.springframework.transaction.annotation.Transactional;

import java.util.List;
import java.util.function.BiConsumer;
import java.util.stream.Stream;

/**
 * @author bjw
 * @since 2024-01-06 11:35
 */
public class ExtendedJpaRepositoryImpl<T, ID> extends SimpleJpaRepository<T, ID> implements ExtendedJpaRepository<T, ID> {
    private final EntityManager entityManager;
    private final JpaEntityInformation<T, ID> entityInformation;
    private final Class<T> classType;

    public ExtendedJpaRepositoryImpl(JpaEntityInformation<T, ID> entityInformation, EntityManager entityManager) {
        super(entityInformation, entityManager);
        this.entityManager = entityManager;
        this.entityInformation = entityInformation;
        this.classType = entityInformation.getJavaType();
    }

    public ExtendedJpaRepositoryImpl(Class<T> domainClass, EntityManager entityManager) {
        super(domainClass, entityManager);
        this.entityManager = entityManager;
        this.classType = domainClass;
        this.entityInformation = null;
    }

    /**
     * Retrieves the identifier of the entity.
     * <p>
     * This method is used to extract the unique identifier (ID) from the given entity object.
     * It is particularly useful in persistence layers, for example, when you need to look up
     * database records based on an entity object. This method overrides the corresponding method
     * in the interface to provide the actual logic for retrieving the identifier.
     *
     * @param entity The entity object from which to extract the identifier.
     * @return The identifier of the entity object.
     */
    @Override
    public ID getId(T entity) {
        return entityInformation.getId(entity);
    }


    @Transactional
    @Override
    public void persist(T entity) {
        entityManager.persist(entity);
    }

    @Transactional
    @Override
    public T merge(T entity) {
        return entityManager.merge(entity);
    }

    @Override
    public List<T> getResultList(BiConsumer<CriteriaBuilder, CriteriaQuery<T>> consumer) {
        return getTypedQuery(consumer).getResultList();
    }

    @Override
    public T getSingleResult(BiConsumer<CriteriaBuilder, CriteriaQuery<T>> consumer) {
        return getTypedQuery(consumer).getSingleResult();
    }

    @Override
    public Stream<T> getResultStream(BiConsumer<CriteriaBuilder, CriteriaQuery<T>> consumer) {
        return getTypedQuery(consumer).getResultStream();
    }

    private TypedQuery<T> getTypedQuery(BiConsumer<CriteriaBuilder, CriteriaQuery<T>> consumer) {

        CriteriaBuilder cb = entityManager.getCriteriaBuilder();

        CriteriaQuery<T> query = cb.createQuery(classType);

        consumer.accept(cb, query);

        return entityManager.createQuery(query);
    }

    /**
     * 执行更新
     *
     * @param consumer 处理函数
     * @return 返回更新的行数
     */
    @Transactional
    @Override
    public int executeUpdate(BiConsumer<CriteriaBuilder, CriteriaUpdate<T>> consumer) {

        CriteriaBuilder cb = entityManager.getCriteriaBuilder();

        CriteriaUpdate<T> query = cb.createCriteriaUpdate(classType);

        consumer.accept(cb, query);

        return entityManager.createQuery(query).executeUpdate();
    }

    /**
     * 执行删除
     *
     * @param consumer 处理函数
     * @return 返回删除的行数
     */
    @Transactional
    @Override
    public int executeDelete(BiConsumer<CriteriaBuilder, CriteriaDelete<T>> consumer) {

        CriteriaBuilder cb = entityManager.getCriteriaBuilder();

        CriteriaDelete<T> query = cb.createCriteriaDelete(classType);

        consumer.accept(cb, query);

        return entityManager.createQuery(query).executeUpdate();
    }

    @Override
    public T findOne(Specification<T> specification, Sort sort) {
        List<T> list = findAll(specification, sort, Limit.of(1));

        return list.isEmpty() ? null : list.get(0);
    }

    @Override
    public List<T> findAll(Specification<T> specification, Sort sort, Limit limit) {
        return this.getQuery(specification, sort).setMaxResults(limit.max()).getResultList();
    }
}
